from functools import lru_cache

import mysql.connector
import numpy as np
from mysql.connector import Error
from omegaconf import OmegaConf


# 创建连接
def create_connection():
    connection = None
    try:
        config = OmegaConf.load("application.yaml")
        connection = mysql.connector.connect(
            host=config.mysql.host,
            user=config.mysql.username,
            password=config.mysql.password,
            database=config.mysql.database,
        )
        print("连接成功")
    except Error as e:
        print(f"连接失败: {e}")
    return connection


# 增加数据
def insert_data(connection, query, data):
    cursor = connection.cursor()
    try:
        cursor.execute(query, data)
        connection.commit()  # 提交事务
        print("数据插入成功")
    except Error as e:
        print(f"数据插入失败: {e}")


# 查询数据
def select_data(connection, query):
    cursor = connection.cursor()
    cursor.execute(query)
    result = cursor.fetchall()  # 获取所有结果
    return result


# 删除数据
def delete_data(connection, query, data):
    cursor = connection.cursor()
    try:
        cursor.execute(query, data)
        connection.commit()  # 提交事务
        print("数据删除成功")
    except Error as e:
        print(f"数据删除失败: {e}")


# 更新数据
def update_data(connection, query, data):
    cursor = connection.cursor()
    try:
        cursor.execute(query, data)
        connection.commit()  # 提交事务
        print("数据更新成功")
    except Error as e:
        print(f"数据更新失败: {e}")


def set_log(connection, train_task_id, link):
    """
    设置训练任务日志链接
    :param connection:
    :param train_task_id:
    :param link:
    :return:
    """
    update_query = "UPDATE train_task SET log=%s WHERE id=%s"
    update_data(connection, update_query, (link, train_task_id))


def update_status(connection, train_task_id, status):
    """
    更新训练状态
    :param connection:
    :param train_task_id:
    :param status:
    :return:
    """
    update_query = "UPDATE train_task SET status=%s WHERE id=%s"
    update_data(connection, update_query, (status, train_task_id))


def set_checkpoint_folder(connection, train_task_id, checkpoint_folder):
    update_query = "UPDATE train_task SET checkpoint_folder=%s WHERE id=%s"
    update_data(connection, update_query, (checkpoint_folder, train_task_id))


def set_other_res(connection, train_task_id, other_res):
    update_query = "UPDATE train_task SET other_res=%s WHERE id=%s"
    update_data(connection, update_query, (other_res, train_task_id))


def set_predict_by_id_batch(connection, predicts):
    """
    predicts = [
    ('value1', 1),
    ('value2', 2),
    ('value3', 3),
    ]
    """
    cursor = connection.cursor()
    sql = "UPDATE inference_item_train_task SET predict=%s WHERE id=%s"
    cursor.executemany(sql, predicts)
    connection.commit()


def select_mask_color_by_dataset_id(connection, dataset_id):
    select_query = f"SELECT color FROM dataset_category WHERE dataset_id={dataset_id} order by 'order' ASC"
    res = select_data(connection, select_query)
    res = [eval(item[0]) for item in res]
    return res


# 主函数
if __name__ == "__main__":
    connection = create_connection()

    # # 示例: 插入数据
    # insert_query = "INSERT INTO your_table (column1, column2) VALUES (%s, %s)"
    # insert_data(connection, insert_query, ("value1", "value2"))
    #
    # # 示例: 查询数据
    # select_query = "SELECT * FROM your_table"
    # results = select_data(connection, select_query)
    # print("查询结果:", results)
    #
    # # 示例: 更新数据
    # update_query = "UPDATE your_table SET column1 = %s WHERE column2 = %s"
    # update_data(connection, update_query, ("new_value1", "value2"))
    #
    # # 示例: 删除数据
    # delete_query = "DELETE FROM your_table WHERE column1 = %s"
    # delete_data(connection, delete_query, ("new_value1",))
    #
    # # 关闭连接
    # if connection.is_connected():
    #     connection.close()
    #     print("连接已关闭")
    select_query = f"""
                select a.name
                    from dataset_category a, train_task b
                where b.id=7 and a.dataset_id = b.dataset_id and a.is_deleted = 0
                order by a.no
                """
    class_name = select_data(connection, select_query)
    class_name = [item[0] for item in class_name]
    print(class_name)
